package rad.framework.jce;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.Key;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.cert.CertificateException;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;

/**
 * The utility is used to create JCE keystore and secret key. Also provides
 * helper method to get the secret key from the keystore.
 * 
 * @author Chunyun Zhao
 */
public class KeyStoreUtil {
	private static final String JCEKS = "jceks";

	private static final String JCE_PROVIDER = "SunJCE";

	private static final String PRNG_ALGORITHM = "SHA1PRNG";

	public static void main(String[] args) throws KeyStoreException,
			NoSuchAlgorithmException, CertificateException, IOException {
		System.out
				.println("The utility is used to create JCE keystore and secret key.");
		System.out
				.println("The keystore file will be created if it doesn't exist. It adds");
		System.out
				.println("the encryption key to the keystore and overwrites existing key ");
		System.out.println("with the same alias.\n");

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		System.out.print("Key store file: ");
		String keyStoreFileName = br.readLine();
		System.out.print("Key store password: ");
		String keyStorePass = br.readLine();
		System.out.print("Key alias: ");
		String keyAlias = br.readLine();
		System.out.print("Key password: ");
		String keyPass = br.readLine();
		System.out.print("Key algorithm (DESede or AES): ");
		String keyAlg = br.readLine();
		if (!"DESede".equals(keyAlg) && !"AES".equals(keyAlg)) {
			throw new RuntimeException(
					"Key algorithm has to be 'DESede' or 'AES'");
		}
		System.out
				.print("Key size (128, 192 or 256 for AES, 112 or 168 for DESede): ");
		int keySize = Integer.parseInt(br.readLine());

		KeyStore keyStore = KeyStore.getInstance(JCEKS);
		SecureRandom seedGenerator = new SecureRandom();
		long seed = seedGenerator.nextLong();

		File file = new File(keyStoreFileName);
		boolean isNewKeyStore;
		if (file.exists()) {
			isNewKeyStore = false;
			System.out.println("Loading from an existing keystore...");
			keyStore.load(new FileInputStream(keyStoreFileName), keyStorePass
					.toCharArray());
		} else {
			isNewKeyStore = true;
			System.out.println("Creating a new keystore...");
			keyStore.load(null, null);
		}
		System.out.println("Generating the secret key...");
		SecretKey secretKey = generateSecretKey(keyAlg, keySize, seed);

		keyStore.setKeyEntry(keyAlias, secretKey, keyPass.toCharArray(), null);
		FileOutputStream fos = new FileOutputStream(keyStoreFileName);
		keyStore.store(fos, keyStorePass.toCharArray());
		System.out.println("The keystore file is "
				+ (isNewKeyStore ? "created" : "updated") + " successfully!!!");
	}

	/**
	 * Generates the private key
	 * 
	 * @param key
	 *            The number used as a seed to the psuedo random number
	 *            generator.
	 * 
	 * @return a SecretKey generated from the seed value passed in.
	 */
	public static SecretKey generateSecretKey(String keyAlg, int keySize,
			long key) {
		SecureRandom prng = null;
		KeyGenerator keyGen = null;
		SecretKey secretKey = null;

		try {
			prng = SecureRandom.getInstance(PRNG_ALGORITHM);
			prng.setSeed(twiddle(key));
			keyGen = KeyGenerator.getInstance(keyAlg, JCE_PROVIDER);
			keyGen.init(keySize);
			keyGen.init(prng);

			secretKey = keyGen.generateKey();
		} catch (Exception e) {
			throw new RuntimeException(
					"Unable to generate " + keyAlg + " key.", e);
		}
		return secretKey;
	}

	public static Key getSecretKey(String keyStoreLocation,
			String keyStorePass, String keyAlias, String keyPass) {
		InputStream keyStoreStream = null;
		try {
			KeyStore keyStore = KeyStore.getInstance(JCEKS);
			keyStoreStream = Thread.currentThread().getContextClassLoader()
					.getResourceAsStream(keyStoreLocation);
			keyStore.load(keyStoreStream, keyStorePass.toCharArray());

			return keyStore.getKey(keyAlias, keyPass.toCharArray());
		} catch (Exception e) {
			throw new RuntimeException("Unable to load key: " + keyAlias
					+ " from keystore: " + keyStoreLocation, e);			
		} finally {
			if (keyStoreStream != null) {
				try {
					keyStoreStream.close();
				} catch (IOException ignoreIt) {
				}
			}
		}
	}

	/**
	 * Obfuscating the randon number generator seed value.
	 * 
	 * @param key
	 *            Randon number generator seed value to obfuscate.
	 * 
	 * @return obfuscated seed value.
	 */
	private static long twiddle(long key) {
		long twiddleBytes = (key % 8) << 3;
		return (key ^ twiddleBytes);
	}
}